from PIL import Image
import cv2 as cv
import numpy as np


def add_alpha():
    img = Image.open('./launch.png')
    img = img.convert('RGBA')

    # r, g, b, a = img.split()
    # m = Image.merge('RGB', (r, g, b))
    # m.save('launch.png')

    x, y = img.size
    for i in range(x):
        for j in range(y):
            color = img.getpixel((i, j))
            color = color[:-1] + (0,)
            img.putpixel((i, j), color)
    img.save('add_alpha.png')


def merge():
    img = cv.imread('./launch.png')
    gray = cv.cvtColor(img, cv.COLOR_BGR2GRAY)
    ret, binary = cv.threshold(gray, 127, 255, cv.THRESH_BINARY_INV)

    # cv.imshow('binary', binary)
    # cv.waitKey(0)
    # cv.destroyAllWindows()

    new_img = np.zeros((img.shape[0], img.shape[1], 4), dtype=np.uint8)
    new_img[:, :, 0] = img[:, :, 0]
    new_img[:, :, 1] = img[:, :, 1]
    new_img[:, :, 2] = img[:, :, 2]
    new_img[:, :, 3] = binary
    cv.imwrite('new_img.png', new_img)


if __name__ == '__main__':
    # add_alpha()
    merge()
